#!/usr/bin/env python3
"""
LLM独立应用主服务器
基于FastAPI的大语言模型对话服务，支持Function Calling和MCP集成
"""

import asyncio
import logging
import signal
import sys
import json
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Optional, Dict, Any, List

import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

# 添加项目根目录到Python路径
sys.path.insert(0, str(Path(__file__).parent))

from config import get_config, AppConfig
from app import LLMClient, ConversationManager
from utils.function_call import FunctionRegistry, get_enabled_schemas
from utils.mcp import get_mcp_client, cleanup_mcp_client, get_enhanced_function_schemas, enhanced_function_call

# 全局变量
llm_client: Optional[LLMClient] = None
conversation_manager: Optional[ConversationManager] = None
app_config: Optional[AppConfig] = None
active_connections: set = set()

# 初始化配置和日志
def init_logging_and_config():
    global app_config
    app_config = get_config()
    
    # 配置日志
    logging.basicConfig(
        level=getattr(logging, app_config.logging.level),
        format=app_config.logging.format
    )

init_logging_and_config()
logger = logging.getLogger("LLM_MAIN")


@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    global llm_client, conversation_manager, app_config

    logger.info("启动LLM独立应用...")

    try:
        # 初始化LLM客户端
        llm_client = LLMClient(app_config)
        logger.info("LLM客户端初始化完成")

        # 初始化对话管理器
        conversation_manager = ConversationManager(app_config)
        logger.info("对话管理器初始化完成")

        # 初始化MCP客户端
        await get_mcp_client()
        logger.info("MCP客户端初始化完成")

        # 导入Function Calling工具（触发注册）
        from utils.calculator.calculator import calculator_func, advanced_calculator_func
        from utils.function_call.music_functions import (
            music_list_func, music_play_func, music_play_random_func, music_stop_func,
            music_volume_set_func, music_volume_up_func, music_volume_down_func,
            music_volume_set_smart_func, music_status_func
        )

        # 初始化音乐播放器
        from utils.function_call import initialize_music_player
        initialize_music_player()
        logger.info("音乐播放器已初始化")

        # 导入云台控制工具（触发注册）
        from utils.function_call.gimbal_functions import (
            gimbal_control_func, gimbal_move_relative_func, gimbal_reset_func,
            gimbal_status_func, gimbal_connect_func, gimbal_disconnect_func
        )

        # 初始化云台控制器
        from utils.function_call.gimbal_functions import initialize_gimbal_controller
        if app_config.gimbal.enabled:
            try:
                initialize_gimbal_controller(
                    port=app_config.gimbal.serial_port,
                    baudrate=app_config.gimbal.baudrate,
                    pan_pin=app_config.gimbal.pan_pin,
                    tilt_pin=app_config.gimbal.tilt_pin
                )
                logger.info("云台控制器已初始化")
            except Exception as e:
                logger.warning(f"云台控制器初始化失败: {e}")

        # 导入VLM工具（触发注册）
        if app_config.vlm.enabled:
            try:
                from utils.function_call.vlm_functions import initialize_vlm_client
                initialize_vlm_client()
                logger.info("VLM客户端已初始化")
            except Exception as e:
                logger.warning(f"VLM客户端初始化失败: {e}")

        # 导入机械臂控制工具（触发注册）
        try:
            logger.info("正在导入机械臂控制工具...")
            from utils.function_call.arm_functions import (
                arm_catch_object_func, arm_put_object_func, arm_stack_put_object_func,
                arm_init_action_func, arm_move_direction_func, arm_gripper_control_func,
                arm_move_to_position_func, arm_status_func
            )
            logger.info("机械臂控制工具导入成功")

            # 初始化机械臂控制器
            from utils.function_call import initialize_arm_controller
            try:
                initialize_arm_controller()
                logger.info("机械臂控制器已初始化")
            except Exception as e:
                logger.warning(f"机械臂控制器初始化失败: {e}")
        except Exception as e:
            logger.error(f"机械臂控制工具加载失败: {e}", exc_info=True)

        logger.info("Function Calling工具已加载")

        yield

    except Exception as e:
        logger.error(f"应用启动失败: {e}", exc_info=True)
        raise
    finally:
        # 清理资源
        logger.info("正在关闭LLM应用...")
        await cleanup_mcp_client()
        logger.info("LLM应用已关闭")


# 创建FastAPI应用
app = FastAPI(
    title="LLM独立应用",
    description="大语言模型对话服务，支持Function Calling和MCP集成",
    version="1.0.0",
    lifespan=lifespan
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 挂载静态文件
app.mount("/static", StaticFiles(directory="static"), name="static")


# 请求/响应模型
class ChatRequest(BaseModel):
    message: str
    session_id: str = "default"
    stream: bool = True


class ChatResponse(BaseModel):
    response: str
    session_id: str
    usage: Optional[Dict[str, Any]] = None


async def broadcast_message(message: dict):
    """广播消息到所有WebSocket连接"""
    global active_connections
    if not active_connections:
        return
    
    message_text = json.dumps(message, ensure_ascii=False)
    disconnected = set()
    
    for websocket in active_connections.copy():
        try:
            await websocket.send_text(message_text)
        except Exception:
            disconnected.add(websocket)
    
    # 清理断开的连接
    active_connections -= disconnected


@app.get("/")
async def root():
    """主页面 - 返回静态HTML文件"""
    return FileResponse("static/index.html")


@app.get("/health")
async def health_check():
    """健康检查"""
    global llm_client, conversation_manager

    try:
        mcp_client = await get_mcp_client()
        mcp_status = mcp_client.get_server_status()
    except Exception as e:
        mcp_status = {"error": str(e)}

    status = {
        "status": "healthy",
        "llm_client": llm_client is not None,
        "conversation_manager": conversation_manager is not None,
        "mcp_client": mcp_status,
        "config_loaded": app_config is not None
    }

    if not (llm_client and conversation_manager and app_config):
        status["status"] = "unhealthy"
        return JSONResponse(status_code=503, content=status)

    return status


@app.get("/config")
async def get_config_info():
    """获取配置信息"""
    global app_config
    
    if not app_config:
        raise HTTPException(status_code=500, detail="配置未加载")
    
    return {
        "model": app_config.llm.model.name,
        "function_calling_enabled": app_config.function_calling.enabled,
        "mcp_enabled": app_config.mcp.enabled,
        "max_history": app_config.conversation.max_history,
        "vlm_enabled": app_config.vlm.enabled
    }


@app.post("/chat", response_model=ChatResponse)
async def chat_endpoint(request: ChatRequest):
    """非流式对话接口"""
    global llm_client, conversation_manager

    if not llm_client or not conversation_manager:
        raise HTTPException(status_code=500, detail="服务未初始化")

    try:
        # 添加用户消息到对话历史
        conversation_manager.add_user_message(request.session_id, request.message)

        # 获取对话消息
        messages = conversation_manager.get_messages_for_api(request.session_id)

        # 获取可用工具
        tools = get_enabled_schemas(app_config)
        mcp_tools = await get_enhanced_function_schemas()
        tools.extend(mcp_tools)

        # 调用LLM
        response = await llm_client.chat(
            messages=messages,
            tools=tools if tools else None,
            stream=False
        )

        # 处理响应
        assistant_message = response['choices'][0]['message']
        content = assistant_message.get('content', '')

        # 处理工具调用
        if 'tool_calls' in assistant_message and assistant_message['tool_calls']:
            tool_results = []
            for tool_call in assistant_message['tool_calls']:
                try:
                    tool_name = tool_call['function']['name']
                    arguments_str = tool_call['function'].get('arguments', '{}')

                    # 安全解析JSON参数
                    if not arguments_str or arguments_str.strip() == '':
                        arguments = {}
                    else:
                        try:
                            arguments = json.loads(arguments_str)
                        except json.JSONDecodeError as e:
                            logger.error(f"工具参数JSON解析失败: {arguments_str}, 错误: {e}")
                            logger.error(f"原始arguments字符串长度: {len(arguments_str)}")
                            logger.error(f"原始arguments内容: {repr(arguments_str)}")
                            arguments = {}

                    logger.info(f"调用工具: {tool_name}, 参数: {arguments}")
                    logger.debug(f"工具调用详情: ID={tool_call.get('id')}, Name={tool_name}, Args={arguments}")

                    # 验证参数完整性
                    if not arguments and tool_name == 'calculator':
                        logger.warning(f"计算器工具缺少必需参数，跳过调用")
                        result = "错误：计算器工具需要操作类型(op)、第一个数字(a)和第二个数字(b)参数"
                    else:
                        # 调用工具 - 使用统一的增强函数调用
                        result = await enhanced_function_call(tool_name, **arguments)

                    logger.info(f"工具调用结果: {result}")
                    tool_results.append(result)

                    # 添加工具调用结果到对话历史
                    conversation_manager.add_tool_message(
                        request.session_id,
                        tool_call.get('id', f"call_{tool_name}"),
                        str(result)
                    )

                except Exception as tool_error:
                    logger.error(f"工具调用失败: {tool_call}, 错误: {tool_error}", exc_info=True)
                    error_result = f"工具调用失败: {str(tool_error)}"
                    tool_results.append(error_result)

                    # 添加错误结果到对话历史
                    conversation_manager.add_tool_message(
                        request.session_id,
                        tool_call.get('id', f"call_error"),
                        error_result
                    )

            # 如果有工具调用，需要再次调用LLM获取最终回复
            if tool_results:
                messages = conversation_manager.get_messages_for_api(request.session_id)
                final_response = await llm_client.chat(
                    messages=messages,
                    stream=False
                )
                content = final_response['choices'][0]['message'].get('content', '')

        # 添加助手回复到对话历史
        conversation_manager.add_assistant_message(request.session_id, content)

        return ChatResponse(
            response=content,
            session_id=request.session_id,
            usage=response.get('usage')
        )

    except Exception as e:
        logger.error(f"对话处理失败: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket端点，用于实时对话"""
    global active_connections, llm_client, conversation_manager

    await websocket.accept()
    active_connections.add(websocket)
    logger.info(f"WebSocket连接已建立，当前连接数: {len(active_connections)}")

    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            message_data = json.loads(data)

            message = message_data.get('message', '')
            image_data = message_data.get('image_data')
            session_id = message_data.get('session_id', 'default')

            # 如果没有文本消息且没有图像数据，则跳过
            if not message and not image_data:
                continue

            # 构建用户消息内容
            user_message_content = message
            if image_data:
                user_message_content += " [用户上传了图像]"

            # 添加用户消息到对话历史
            conversation_manager.add_user_message(session_id, user_message_content)

            # 如果有图像数据且启用了VLM功能，自动调用图像分析
            vlm_result = None
            if image_data and app_config.vlm.enabled and FunctionRegistry.is_registered('vlm_analyze_image'):
                try:
                    # 调用VLM图像分析功能
                    vlm_result = await FunctionRegistry.call('vlm_analyze_image', image_base64=image_data)
                    logger.info(f"VLM图像分析结果: {vlm_result}")

                    # 将分析结果添加到对话历史中
                    if vlm_result.get('success'):
                        analysis_content = f"图像分析结果: {vlm_result.get('description', '无描述')}"
                        conversation_manager.add_user_message(session_id, analysis_content)
                    else:
                        error_content = f"图像分析失败: {vlm_result.get('error', '未知错误')}"
                        conversation_manager.add_user_message(session_id, error_content)
                except Exception as e:
                    logger.error(f"VLM图像分析调用失败: {e}")
                    conversation_manager.add_user_message(session_id, f"图像分析调用失败: {str(e)}")

            logger.info(f"用户消息: {message}, 会话ID: {session_id}")

            # 实现多轮工具调用循环
            MAX_TOOL_CALL_ROUNDS = 5
            tool_call_round_count = 0

            while tool_call_round_count < MAX_TOOL_CALL_ROUNDS:
                logger.info(f"--- 多轮工具调用：第 {tool_call_round_count + 1} 轮 ---")

                # 获取对话消息
                messages = conversation_manager.get_messages_for_api(session_id)

                # 获取可用工具
                tools = get_enabled_schemas(app_config)
                mcp_tools = await get_enhanced_function_schemas()
                tools.extend(mcp_tools)

                # 流式响应
                full_content = ""
                tool_calls_buffer = {}  # 工具调用缓冲区，用于拼接分片的arguments

                async for chunk in llm_client.stream_chat(
                    messages=messages,
                    tools=tools if tools else None,
                    stream=True
                ):
                    if "error" in chunk:
                        await websocket.send_text(json.dumps({
                            "type": "error",
                            "error": chunk["error"],
                            "hint": chunk["hint"]
                        }, ensure_ascii=False))
                        break

                    # 处理内容
                    if "content" in chunk:
                        content = chunk["content"]
                        full_content += content
                        await websocket.send_text(json.dumps({
                            "type": "content",
                            "content": content,
                            "session_id": session_id
                        }, ensure_ascii=False))

                    # 处理工具调用 - 支持流式分片arguments拼接
                    if "tool_calls" in chunk and chunk["tool_calls"]:
                        for tool_call in chunk["tool_calls"]:
                            if not isinstance(tool_call, dict) or 'function' not in tool_call:
                                logger.warning(f"跳过无效的工具调用数据: {tool_call}")
                                continue

                        # 获取工具调用索引，用于缓冲区管理
                        idx = tool_call.get("index", 0)

                        # 深拷贝以避免修改原始chunk
                        tool_call_copy = json.loads(json.dumps(tool_call))

                        if idx not in tool_calls_buffer:
                            # 新的工具调用，初始化缓冲区
                            tool_calls_buffer[idx] = tool_call_copy
                            # 确保必要字段存在
                            if 'id' not in tool_calls_buffer[idx]:
                                tool_calls_buffer[idx]['id'] = f"call_{tool_call_copy['function'].get('name', 'unknown')}_{idx}"
                            if 'arguments' not in tool_calls_buffer[idx]['function']:
                                tool_calls_buffer[idx]['function']['arguments'] = ""
                        else:
                            # 已存在的工具调用，拼接arguments
                            existing_args = tool_calls_buffer[idx]["function"].get("arguments", "")
                            new_args_piece = tool_call_copy["function"].get("arguments", "")
                            tool_calls_buffer[idx]["function"]["arguments"] = existing_args + new_args_piece

                # 处理结束
                if "finish_reason" in chunk:
                    # 将缓冲区中的工具调用转换为最终列表
                    tool_calls = [tool_calls_buffer[idx] for idx in sorted(tool_calls_buffer.keys())]

                    # 如果有工具调用，执行工具
                    if tool_calls:
                        await websocket.send_text(json.dumps({
                            "type": "tool_calling",
                            "message": "正在调用工具..."
                        }, ensure_ascii=False))

                        # 添加助手消息（包含工具调用）
                        conversation_manager.add_assistant_message(
                            session_id,
                            full_content,
                            tool_calls=tool_calls
                        )

                        # 执行工具调用
                        for tool_call in tool_calls:
                            tool_name = "未知工具"  # 默认值
                            try:
                                tool_name = tool_call['function']['name']
                                arguments_str = tool_call['function'].get('arguments', '{}')

                                # 发送开始调用特定工具的消息
                                await websocket.send_text(json.dumps({
                                    "type": "tool_start",
                                    "tool_name": tool_name,
                                    "tool_id": tool_call.get('id', f"call_{tool_name}")
                                }, ensure_ascii=False))

                                # 安全解析JSON参数
                                if not arguments_str or arguments_str.strip() == '':
                                    arguments = {}
                                else:
                                    try:
                                        arguments = json.loads(arguments_str)
                                    except json.JSONDecodeError as e:
                                        logger.error(f"工具参数JSON解析失败: {arguments_str}, 错误: {e}")
                                        # 如果解析失败，记录原始字符串以便调试
                                        logger.error(f"原始arguments字符串长度: {len(arguments_str)}")
                                        logger.error(f"原始arguments内容: {repr(arguments_str)}")
                                        arguments = {}

                                logger.info(f"调用工具: {tool_name}, 参数: {arguments}")
                                logger.debug(f"工具调用详情: ID={tool_call.get('id')}, Name={tool_name}, Args={arguments}")

                                # 验证参数完整性
                                if not arguments and tool_name == 'calculator':
                                    logger.warning(f"计算器工具缺少必需参数，跳过调用")
                                    result = "错误：计算器工具需要操作类型(op)、第一个数字(a)和第二个数字(b)参数"
                                else:
                                    # 调用工具 - 使用统一的增强函数调用
                                    result = await enhanced_function_call(tool_name, **arguments)

                                logger.info(f"工具调用结果: {result}")

                                # 发送工具调用结果
                                await websocket.send_text(json.dumps({
                                    "type": "tool_result",
                                    "tool_name": tool_name,
                                    "tool_id": tool_call.get('id', f"call_{tool_name}"),
                                    "result": str(result)
                                }, ensure_ascii=False))

                                # 添加工具结果到对话历史
                                conversation_manager.add_tool_message(
                                    session_id,
                                    tool_call.get('id', f"call_{tool_name}"),
                                    str(result)
                                )

                            except Exception as tool_error:
                                logger.error(f"工具调用失败: {tool_call}, 错误: {tool_error}", exc_info=True)

                                # 发送工具错误结果
                                await websocket.send_text(json.dumps({
                                    "type": "tool_result",
                                    "tool_name": tool_name,
                                    "tool_id": tool_call.get('id', f"call_error"),
                                    "result": f"工具调用失败: {str(tool_error)}",
                                    "error": True
                                }, ensure_ascii=False))

                                # 添加错误结果到对话历史
                                conversation_manager.add_tool_message(
                                    session_id,
                                    tool_call.get('id', f"call_error"),
                                    f"工具调用失败: {str(tool_error)}"
                                )

                        # 工具调用完成，增加轮次计数，继续下一轮
                        tool_call_round_count += 1
                        logger.info(f"工具调用完成，准备进入第 {tool_call_round_count + 1} 轮")
                        # 继续while循环，不break
                    else:
                        # 没有工具调用，直接添加回复并结束循环
                        conversation_manager.add_assistant_message(session_id, full_content)

                        await websocket.send_text(json.dumps({
                            "type": "complete",
                            "session_id": session_id
                        }, ensure_ascii=False))

                        # 结束多轮工具调用循环
                        logger.info("没有工具调用，多轮工具调用循环结束")
                        break  # 跳出while循环
                    # 工具调用情况下，继续while循环

            # 检查是否超过最大轮数
            if tool_call_round_count >= MAX_TOOL_CALL_ROUNDS:
                logger.info("达到最大工具调用轮数，循环结束")
                break

    except WebSocketDisconnect:
        logger.info("WebSocket客户端断开连接")
    except Exception as e:
        logger.error(f"WebSocket错误: {e}", exc_info=True)
    finally:
        active_connections.discard(websocket)
        logger.info(f"WebSocket连接已断开，当前连接数: {len(active_connections)}")


@app.get("/conversations/{session_id}")
async def get_conversation(session_id: str):
    """获取对话历史"""
    global conversation_manager

    if not conversation_manager:
        raise HTTPException(status_code=500, detail="对话管理器未初始化")

    conversation = conversation_manager.get_conversation(session_id)
    summary = conversation_manager.get_conversation_summary(session_id)

    return {
        "session_id": session_id,
        "messages": conversation,
        "summary": summary
    }


@app.delete("/conversations/{session_id}")
async def clear_conversation(session_id: str):
    """清空对话历史"""
    global conversation_manager

    if not conversation_manager:
        raise HTTPException(status_code=500, detail="对话管理器未初始化")

    conversation_manager.clear_conversation(session_id)

    return {"message": f"会话 {session_id} 已清空"}


@app.get("/tools")
async def get_available_tools():
    """获取可用工具列表"""
    from utils.mcp.client.mcp_client import get_mcp_tools_schemas

    tools = {
        "function_calling": get_enabled_schemas(app_config),
        "mcp_tools": await get_mcp_tools_schemas(),
        "registered_functions": FunctionRegistry.list_functions()
    }

    return tools


@app.get("/mcp-config")
async def get_mcp_config():
    """获取MCP服务器配置信息"""
    import json
    from pathlib import Path

    config_path = Path("utils/mcp/config/mcp_servers.json")
    if not config_path.exists():
        return {"mcpServers": {}}

    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        return config
    except Exception as e:
        logger.error(f"读取MCP配置文件失败: {e}")
        return {"mcpServers": {}}


def signal_handler(signum, frame):
    """信号处理器"""
    logger.info(f"收到信号 {signum}，正在关闭应用...")
    sys.exit(0)


def main():
    """主函数"""
    global app_config

    # 设置信号处理器
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    # 启动服务器
    uvicorn.run(
        "main:app",
        host=app_config.server.host,
        port=app_config.server.port,
        reload=app_config.server.debug,
        log_level=app_config.logging.level.lower()
    )


if __name__ == "__main__":
    main()
